import numpy as np
import torch
from scipy.interpolate import interp1d
from statsmodels.tsa.stattools import adfuller
import pandas as pd
import os
import hashlib
from typing import Callable


def train_test_split_by_ratio(inputs, labels, train_ratio=0.8):
    total_samples = inputs.shape[0]
    train_size = int(total_samples * train_ratio)
    train_inputs = inputs[:train_size]
    train_labels = labels[:train_size]
    test_inputs = inputs[train_size:]
    test_labels = labels[train_size:]
    return train_inputs, train_labels, test_inputs, test_labels


def normalize_features_timewise(x):
    """
    对输入样本 (n, l, c) 的特征在时间维度上进行 z-score 标准化。
    对每只股票、每个特征独立归一化，即：
    对 x[i, :, j] 做 (x - mean) / std 操作。

    若某个时间序列为常数（std = 0），则归一化结果设为 0。
    """
    x = x.copy()
    mean = x.mean(axis=1, keepdims=True)  # 在时间维度上求均值
    std = x.std(axis=1, keepdims=True)  # 在时间维度上求标准差

    # 避免除零错误：将常数列的 std 设置为 1（后面再手动置 0）
    normalized = (x - mean) / (std + 1e-6)

    # 将 std == 0 的序列置为 0
    constant_mask = std < 1e-12
    normalized[constant_mask.repeat(x.shape[1], axis=1)] = 0

    return normalized


def normalize_features(x):
    """
    对输入样本的所有特征都进行 z-score 标准化
    自动处理可能全为常数的特征列（将其置为0）
    """
    x = x.copy()
    mean = x.mean(axis=0, keepdims=True)
    std = x.std(axis=0, keepdims=True)
    constant_columns = std == 0
    x = (x - mean) / (std + 1e-6)  # 正常列
    x[:, constant_columns.flatten()] = 0
    return x


def normalize_features_with_volatility(x):
    """
    对输入样本的所有特征进行 z-score 标准化，并返回第一个特征的波动性。
    Args:
        x (np.ndarray): 输入的样本数据，形状为 (window_length, num_features)。
                        例如 (8, 6) 表示8个时间步，每个时间步6个特征。
    Returns:
        tuple:
            - np.ndarray: 归一化后的特征矩阵，形状与输入x相同。
            - float: 第一个特征（通常是价格）在归一化前的标准差，作为波动性指标。
    """
    x_copy = x.copy()
    volatility = np.std(x_copy[:, 0], ddof=1)

    # --- 2. 对所有特征进行 Z-score 标准化 ---
    mean = np.mean(x_copy, axis=0, keepdims=True)
    std = np.std(x_copy, axis=0, keepdims=True, ddof=1)
    constant_columns = (std == 0).flatten()
    normalized_x = (x_copy - mean) / (std + 1e-8)
    normalized_x[:, constant_columns] = 0
    return normalized_x, volatility


def add_jittering(x, sigma=0.01, num_raw=8):
    """
    只对前 num_raw 个通道加高斯扰动，技术指标通道不加扰动。
    :param x: (T, C)
    :return: 同形状扰动后的数据
    """
    noise = np.zeros_like(x)
    noise[:, :num_raw] = np.random.normal(
        loc=0.0, scale=sigma, size=(x.shape[0], num_raw)
    )
    return x + noise


def time_mask(x, p=0.1):
    x = x.copy()
    num_mask = int(p * x.shape[0])
    mask_indices = np.random.choice(x.shape[0], num_mask, replace=False)
    x[mask_indices, :] = 0  # 或者 np.nan
    return x


def time_shift(x, shift_range=3):
    shift = np.random.randint(-shift_range, shift_range + 1)
    if shift == 0:
        return x
    return np.roll(x, shift, axis=0)


def time_warp(x, sigma=0.2, knot=4):
    orig_steps = np.arange(x.shape[0])  # 原始时间轴
    warp_steps = np.linspace(0, x.shape[0] - 1, knot + 2)  # 控制点（均匀分布）
    warp_factors = np.random.normal(1.0, sigma, knot + 2)  # 随机扰动比例
    warped = np.cumsum(warp_factors)  # 生成非均匀时间轴
    warped = np.interp(orig_steps, np.linspace(0, x.shape[0] - 1, len(warped)), warped)

    x_warped = np.empty_like(x)
    for dim in range(x.shape[1]):
        f = interp1d(warped, x[:, dim], kind="linear", fill_value="extrapolate")
        x_warped[:, dim] = f(orig_steps)
    return x_warped


def multi_diff(x, num_price_cols=4, max_diff_order=3, adf_p_threshold=0.05):
    """
    对前 num_price_cols 列做差分处理，使用 ADF 检验递归判断是否平稳，直到通过或达到 max_diff_order。

    :param x: 输入数据，shape = (T, D)
    :param num_price_cols: 要进行差分的前几列（如开高低收）
    :param max_diff_order: 最大差分阶数，默认允许 3 阶
    :param adf_p_threshold: ADF 检验的显著性阈值（默认 0.05）
    :return: 差分后的数据，shape = (T, D)，前 num_price_cols 可能已被覆盖
    """
    x = x.copy()
    T, D = x.shape
    target = x[:, :num_price_cols]  # (T, N)
    diffed = target.copy()

    for order in range(1, max_diff_order + 1):
        # 对每列执行 ADF 检验，找出不平稳列索引
        p_values = np.array([adfuller(diffed[:, i])[1] for i in range(num_price_cols)])
        unstable_mask = p_values > adf_p_threshold

        if not unstable_mask.any():
            # print(f"[ADF] 所有列在第 {order - 1} 阶差分后已平稳。")
            break

        # 只对不平稳列执行差分
        diffed[:, unstable_mask] = np.diff(
            diffed[:, unstable_mask], axis=0, prepend=diffed[0:1, unstable_mask]
        )
        # print(
        #     f"[ADF] 第 {order} 阶差分，仍不平稳的列索引：{np.where(unstable_mask)[0].tolist()}"
        # )

    # 覆盖原数据中前 num_price_cols
    x[:, :num_price_cols] = diffed
    return x


def process_blc_features(blc_tensor: torch.Tensor, num_price_cols: int = 5):
    """
    对 BLC 格式的状态进行验证集级别的预处理：差分 + 归一化
    :param blc_tensor: shape = (B, L, C)
    :param num_price_cols: 前多少列视为价格类特征，需要做差分（例如开高低收量）
    :return: 同样 shape 的处理后张量
    """
    # 保证是 numpy 操作，复用已有工具
    blc_np = blc_tensor.cpu().numpy()
    B, L, C = blc_np.shape

    processed = []
    for b in range(B):
        x = blc_np[b]  # shape: (L, C)
        x = multi_diff(x, num_price_cols=num_price_cols)
        x = normalize_features(x)
        processed.append(x)

    processed_np = np.stack(processed)  # (B, L, C)
    return torch.tensor(processed_np, dtype=torch.float32, device=blc_tensor.device)


def process_state_to_blc(state, window, future_count, tech_num):
    """
    将 state 转换为 (future_num, window, 1 + tech_num) 的格式
    确保正确处理 tech_ary 的 [len × (future_num*tech_num)] 存储结构

    Args:
        state: 输入状态 [available_amount, future_positions, price_window, tech_window]
        window: 时间窗口长度
        future_count: 期货品种数量
        tech_num: 每个期货的技术指标数量

    Returns:
        torch.Tensor: (future_num, window, 1 + tech_num)
    """
    expected_len = (
        1 + future_count + window * future_count + window * future_count * tech_num
    )
    if len(state) != expected_len:
        raise ValueError(f"State长度错误，预期{expected_len}，实际{len(state)}")
    ptr = 0
    ptr += 1
    future_positions = state[ptr : ptr + future_count]  # future_count个值
    ptr += future_count
    price_size = window * future_count
    price_flat = state[ptr : ptr + price_size]
    ptr += price_size
    price = price_flat.reshape(window, future_count).T
    tech_flat = state[ptr:]
    tech = tech_flat.reshape(window, future_count, tech_num).transpose(1, 0, 2)
    price = price[:, :, np.newaxis]  # (future_count, window, 1)
    state_blc = np.concatenate(
        [price, tech], axis=-1
    )  # (future_count, window, 1+tech_num)
    return state_blc


def dtw_distance(series_a, series_b):
    """
    参数:
    series_a (np.ndarray): 第一个时间序列，形状为 (L, C)。
    series_b (np.ndarray): 第二个时间序列，形状为 (L, C)。
    返回:
    float: 两个序列之间的DTW距离。
    """

    if not isinstance(series_a, np.ndarray) or not isinstance(series_b, np.ndarray):
        raise TypeError("输入必须是 NumPy 数组。")
    if series_a.ndim != 2 or series_b.ndim != 2:
        raise ValueError("输入数组必须是二维的 (L, C)。")
    if series_a.shape[1] != series_b.shape[1]:
        raise ValueError("两个序列的特征维度 C 必须相同。")

    L_a, C = series_a.shape
    L_b, _ = series_b.shape

    cost_matrix = np.full((L_a + 1, L_b + 1), np.inf)
    cost_matrix[0, 0] = 0

    for i in range(1, L_a + 1):
        for j in range(1, L_b + 1):
            cost = np.linalg.norm(series_a[i - 1] - series_b[j - 1])
            last_min_cost = np.min(
                [
                    cost_matrix[i - 1, j],  # 来自下方
                    cost_matrix[i, j - 1],  # 来自左方
                    cost_matrix[i - 1, j - 1],
                ]
            )  # 来自左下角
            cost_matrix[i, j] = cost + last_min_cost

    return cost_matrix[L_a, L_b]


def prepare_sliding_window(
    data,
    lags,
    trend_horizon=1,
    use_log_returns=True,
    clip_range=(-2.0, 2.0),
    augment=False,
):
    """
    一个简化的滑动窗口函数，标签Y直接是未来收益率，无需波动率标准化。

    Args:
        data (np.array): 输入数据，形状为 (b, L, c)
                         b: 资产数量 (batch/assets)
                         L: 时间序列总长度 (Length)
                         c: 特征数量 (channels), 约定 data[:, :, 0] 是价格
        lags (int): 输入窗口的长度 (lookback window size)。
        trend_horizon (int): 预测未来的时间跨度。
        use_log_returns (bool): 是否使用对数收益率。
        clip_range (tuple or None): 对原始收益率进行截断的范围。如果为None，则不截断。
        augment (bool): 是否对输入数据进行增强。

    Returns:
        torch.Tensor: 输入数据 X，形状 (num_windows, b, lags, c)
        torch.Tensor: 标签 Y，形状 (num_windows, b, 1)
    """
    b, L, c = data.shape
    X_list, Y_list = [], []
    epsilon = 1e-9

    # 遍历所有可能的起始点来创建窗口
    for start in range(L - lags - trend_horizon + 1):
        x_batch, y_batch = [], []

        # 遍历每个资产
        for asset_idx in range(b):
            asset_data = data[asset_idx]

            # --- 1. 计算未来的真实收益率 (这是我们现在的核心目标) ---
            # 当前窗口的最后一个价格
            price_now = asset_data[start + lags - 1, 0]
            # 预测时间范围内的平均未来价格
            mean_future_price = np.mean(
                asset_data[start + lags : start + lags + trend_horizon, 0]
            )

            # 计算原始收益率 (raw_return)
            if use_log_returns:
                raw_return = np.log(mean_future_price / (price_now + epsilon))
            else:
                raw_return = (mean_future_price - price_now) / (price_now + epsilon)

            # --- 2. (可选) 对原始收益率进行截断 ---
            if clip_range is not None:
                label = np.clip(raw_return, clip_range[0], clip_range[1])
            else:
                label = raw_return

            # --- 3. 准备输入序列 X ---
            x_seq = asset_data[start : start + lags].copy()

            # 假设有一个特征归一化函数
            # x_seq = normalize_features(x_seq)
            if augment:
                x_seq = normalize_features(x_seq)
                x_seq = time_mask(x_seq, p=0.15)
                x_seq = time_warp(x_seq, sigma=0.2, knot=6)
                x_seq = add_jittering(x_seq, sigma=0.01)
            else:
                x_seq = normalize_features(x_seq)

            x_batch.append(x_seq)
            y_batch.append([label])  # 标签Y现在就是简化的收益率

        X_list.append(x_batch)
        Y_list.append(y_batch)

    # 将列表转换为Numpy数组，然后再转换为PyTorch张量
    X = torch.tensor(np.array(X_list), dtype=torch.float32)
    Y = torch.tensor(np.array(Y_list), dtype=torch.float32)

    return X, Y


def process_futures_data_blc(
    futures_list, folder_path, cache_dir="./cache", use_cache=True
):
    """
    处理期货数据，返回一个 (b, l, c) 形状的 NumPy 数组。
    新增了缓存功能，可大幅提升重复调用的速度。

    B: Batch size, 即期货合约的数量
    L: Sequence length, 即对齐后的公共时间戳数量
    C: Channels, 即每个合约的特征数量

    参数:
    futures_list (list): 期货合约的名称列表, 例如 ['a', 'ag', 'al']。
    folder_path (str): 存放CSV文件的文件夹路径。
    cache_dir (str, optional): 存放缓存文件的文件夹。默认为 "./cache"。
    use_cache (bool, optional): 是否启用缓存功能。默认为 True。

    返回:
    numpy.ndarray: 形状为 (b, l, c) 的 NumPy 数组，如果出错则返回 None。
    """
    # --- 步骤 1: 检查缓存 ---
    if use_cache:
        # a. 根据输入列表生成一个确定性的文件名
        # 先排序，确保 ['a', 'b'] 和 ['b', 'a'] 生成相同的文件名
        sorted_futures = sorted(futures_list)
        identifier = "_".join(sorted_futures).encode("utf-8")

        # 使用md5哈希生成一个简短且唯一的文件名
        hash_name = hashlib.md5(identifier).hexdigest()
        cache_filename = f"{hash_name}.npy"
        cache_path = os.path.join(cache_dir, cache_filename)

        # b. 检查缓存文件是否存在
        if os.path.exists(cache_path):
            print(f"发现缓存，从 {cache_path} 加载数据...")
            try:
                cached_data = np.load(cache_path)
                return cached_data
            except Exception as e:
                print(f"加载缓存文件时出错: {e}。将重新进行数据处理。")

    # --- 缓存未命中或未使用缓存，执行原始处理流程 ---
    print("未找到有效缓存或已禁用缓存，开始从原始文件处理数据...")
    processed_dfs = []

    for future in futures_list:
        file_path = os.path.join(folder_path, f"{future}_tech.csv")
        # ... (后续的文件读取和处理逻辑与您原始代码完全相同)
        if not os.path.exists(file_path):
            print(f"警告: 文件 {file_path} 不存在，已跳过。")
            continue
        try:
            df = pd.read_csv(file_path, index_col=0)
            df.index = pd.to_datetime(df.index)

            # --- 核心修改在这里 ---
            # 1. (可选但推荐) 替换无穷大的值为 NaN
            df.replace([np.inf, -np.inf], np.nan, inplace=True)

            # 2. (关键步骤) 删除任何包含 NaN 值的行
            original_len = len(df)
            df.dropna(inplace=True)
            if len(df) < original_len:
                print(
                    f"在 {file_path} 中删除了 {original_len - len(df)} 行包含 NaN 的数据。"
                )

            if df.empty:
                print(f"警告: 文件 {file_path} 在清理 NaN 后变为空，已跳过。")
                continue

            processed_dfs.append(df)
        except Exception as e:
            print(f"处理文件 {file_path} 时发生错误: {e}")
            continue

    if not processed_dfs:
        print("没有成功处理任何文件。")
        return None

    num_features = processed_dfs[0].shape[1]
    if not all(df.shape[1] == num_features for df in processed_dfs):
        print("错误: 输入的CSV文件特征数量（列数）不一致。")
        return None

    common_index = processed_dfs[0].index
    for df in processed_dfs[1:]:
        common_index = common_index.intersection(df.index)

    if common_index.empty:
        print("警告: 在所有文件中未找到任何公共时间戳。")
        return np.array([]).reshape((len(processed_dfs), 0, num_features))

    aligned_arrays = [df.loc[common_index].to_numpy() for df in processed_dfs]
    final_array = np.stack(aligned_arrays, axis=0)

    # --- 步骤 2: 保存结果到缓存 ---
    if use_cache:
        # a. 确保缓存目录存在
        os.makedirs(cache_dir, exist_ok=True)

        # b. 保存NumPy数组
        try:
            print(f"处理完成，将结果保存到缓存文件: {cache_path}")
            np.save(cache_path, final_array)
        except Exception as e:
            print(f"保存缓存文件时出错: {e}")

    return final_array


def calculate_dtw_adjacency_mask_optimized(x_sample, r=0.2):
    """
    优化版本：只进行一次 .cpu().numpy() 转换
    """
    Stock_Num, L, C = x_sample.shape
    adj = torch.zeros(Stock_Num, Stock_Num, device=x_sample.device)

    # !! 性能优化 !!
    # 一次性将整个样本移至CPU并转为Numpy，避免在循环中反复操作
    x_sample_np = x_sample.cpu().numpy()

    for i in range(Stock_Num):
        for j in range(i, Stock_Num):
            # 现在直接使用Numpy数组进行计算
            dist = dtw_distance(x_sample_np[i], x_sample_np[j])
            similarity = 1 / (1 + dist)
            adj[i, j] = similarity
            adj[j, i] = similarity

    adj.fill_diagonal_(0.0)

    num_edges = Stock_Num * Stock_Num
    k = int(num_edges * r)

    if k == 0:
        mask = torch.zeros_like(adj)
    else:
        values = adj.flatten()
        threshold = torch.topk(values, k, sorted=True).values[-1]
        mask = (adj >= threshold).float()

    mask.fill_diagonal_(1.0)
    return mask


def batch_calculate_dtw_adjacency_mask(x_batch, r=0.2):
    """
    调用优化后的函数并返回一个单一的三维张量
    """
    # 使用列表推导式让代码更简洁
    result_list = [
        calculate_dtw_adjacency_mask_optimized(x_batch[i], r=r)
        for i in range(x_batch.shape[0])
    ]

    # 同样使用 torch.stack 进行堆叠
    return torch.stack(result_list, dim=0)


def load_or_create_dataset(
    futures_list,
    data_folder,
    lags=32,
    cache_dir="./cache",
    use_cache=True,
    val_ratio=0.2,  # 验证集比例
):
    """
    加载或生成期货数据集，返回训练集和验证集的数据与邻接矩阵。

    参数:
        futures_list (list): 期货合约列表。
        data_folder (str): 原始数据文件夹。
        lags (int): 滑动窗口长度。
        cache_dir (str): 缓存目录。
        use_cache (bool): 是否启用缓存。
        val_ratio (float): 验证集比例。

    返回:
        (X_train, y_train, matrices_train,
         X_val, y_val, matrices_val)
    """
    # === 1. 缓存处理 ===
    if use_cache:
        sorted_futures_str = "_".join(sorted(futures_list))
        identifier = f"dataset_{sorted_futures_str}_lags{lags}_val{val_ratio}".encode(
            "utf-8"
        )
        hash_name = hashlib.md5(identifier).hexdigest()
        cache_filename = f"{hash_name}.npz"
        cache_path = os.path.join(cache_dir, cache_filename)

        if os.path.exists(cache_path):
            print(f"发现最终数据集缓存，从 {cache_path} 加载...")
            try:
                data = np.load(cache_path)
                return (
                    data["x_train"],
                    data["y_train"],
                    data["matrices_train"],
                    data["x_val"],
                    data["y_val"],
                    data["matrices_val"],
                )
            except Exception as e:
                print(f"加载缓存文件失败: {e}。将重新生成数据。")

    # === 2. 加载原始数据 ===
    print("=" * 30)
    print("未找到最终数据集缓存，开始完整数据处理流程...")
    initial_data = process_futures_data_blc(
        futures_list, data_folder, cache_dir=cache_dir, use_cache=use_cache
    )
    if initial_data is None:
        print("错误：加载原始数据失败。")
        return None, None, None, None, None, None

    # === 3. 创建滑动窗口 ===
    # 训练集: augment=True
    X_all, y_all = prepare_sliding_window(initial_data, lags=lags, augment=False)

    # 切分训练集和验证集
    total_samples = X_all.shape[0]
    split_idx = int(total_samples * (1 - val_ratio))
    y_train = y_all[:split_idx]
    X_val, y_val = X_all[split_idx:], y_all[split_idx:]

    # 对训练集再生成增强版
    X_train, _ = prepare_sliding_window(initial_data, lags=lags, augment=True)
    X_train = X_train[:split_idx]  # 对齐切分后的索引

    print(
        f"步骤 1/2 完成: "
        f"X_train {X_train.shape}, y_train {y_train.shape} "
        f"X_val {X_val.shape}, y_val {y_val.shape}"
    )

    # === 4. 计算邻接矩阵 ===
    matrices_train = batch_calculate_dtw_adjacency_mask(X_train)
    matrices_val = batch_calculate_dtw_adjacency_mask(X_val)
    print(
        f"步骤 2/2 完成: "
        f"matrices_train {matrices_train.shape}, "
        f"matrices_val {matrices_val.shape}"
    )
    print("=" * 30)

    # === 5. 保存缓存 ===
    if use_cache:
        os.makedirs(cache_dir, exist_ok=True)
        print(f"处理完成，将最终数据集保存到缓存: {cache_path}")
        try:
            np.savez_compressed(
                cache_path,
                x_train=X_train,
                y_train=y_train,
                matrices_train=matrices_train,
                x_val=X_val,
                y_val=y_val,
                matrices_val=matrices_val,
            )
        except Exception as e:
            print(f"保存缓存文件失败: {e}")

    return (
        X_train,
        y_train,
        matrices_train,
        X_val,
        y_val,
        matrices_val,
    )
